import os

from flask import Flask, render_template, request

from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.embeddings import QianfanEmbeddingsEndpoint
from langchain.vectorstores import Chroma
from langchain.chat_models import QianfanChatEndpoint
from langchain.prompts import PromptTemplate
from langchain.memory import ConversationSummaryMemory
from langchain.chains import ConversationalRetrievalChain
from langchain.document_loaders import PyPDFLoader
from langchain.document_loaders import UnstructuredFileLoader

os.environ['QIANFAN_AK'] = "a6OGDgMKaWCHpiUP3afGfR9U"
os.environ['QIANFAN_SK'] = "ijE4BGBxsNbv689h3j04ko3zzPL6SVM0"

CUSTOM_PROMPT_TEMPLATE = """
    使用下面的语料来回答本模板最末尾的问题。如果你不知道问题的答案，直接回答 "我不知道"，禁止随意编造答案。
    为了保证答案尽可能简洁，你的回答必须不超过三句话。
    请注意！在每次回答结束之后，你都必须接上 "感谢你的提问" 作为结束语
    以下是一对问题和答案的样例：
        请问：秦始皇的原名是什么
        秦始皇原名嬴政。感谢你的提问。

    以下是语料：
    
    {context}
    
    请问：{question}
"""

loader = PyPDFLoader("采购计划 产品操作说明书.pdf")
data = loader.load()

text_splitter = RecursiveCharacterTextSplitter(chunk_size = 384, chunk_overlap = 0, separators=["\n\n", "\n", " ", "", "。", "，"])
all_splits = text_splitter.split_documents(data)

vectorstore = Chroma.from_documents(documents=all_splits, embedding=QianfanEmbeddingsEndpoint())

QA_CHAIN_PROMPT = PromptTemplate.from_template(CUSTOM_PROMPT_TEMPLATE)

llm = QianfanChatEndpoint(streaming=True)
retriever=vectorstore.as_retriever(search_type="similarity_score_threshold", search_kwargs={'score_threshold': 0.0})

memory = ConversationSummaryMemory(llm=llm,memory_key="chat_history",return_messages=True)
qa = ConversationalRetrievalChain.from_llm(llm, retriever=retriever, memory=memory, combine_docs_chain_kwargs={"prompt": QA_CHAIN_PROMPT})

app = Flask(__name__)

@app.route('/')
def home():
    return 'home page'

@app.route('/question')
def about():
    question = request.args['question']
    result = qa(question)
    return result['answer']

if __name__ == '__main__':
    app.run(debug=True, host="localhost", port=5000)
